<#
.SYNOPSIS
SSH Key Auto Configuration Script
.DESCRIPTION
Automatically install OpenSSH and use PowerShell native methods to achieve fully automated SSH key configuration for multiple servers
.AUTHOR
Ansible Project Team
.DATE
2024-01-20
.VERSION
2.0 - Using PowerShell native methods instead of sshpass
#>

# Set console encoding to UTF-8 to solve Chinese display garbled characters
try {
    [Console]::OutputEncoding = [System.Text.Encoding]::UTF8
    $OutputEncoding = [System.Text.Encoding]::UTF8
    # Force set console code page to UTF-8
    cmd /c "chcp 65001 >nul 2>&1"
} catch {
    # Ignore encoding errors
}
# Set PowerShell session output encoding
$PSDefaultParameterValues['Out-File:Encoding'] = 'utf8'
$PSDefaultParameterValues['*:Encoding'] = 'utf8'

# Configuration parameters
$servers = @('192.168.144.132', '192.168.144.133', '192.168.144.135')
$username = 'zenglg'
$password = 'Tuners2012'
$keyPath = "$env:USERPROFILE\.ssh\id_rsa"
$pubKeyPath = "$env:USERPROFILE\.ssh\id_rsa.pub"

# Function: Install OpenSSH Client and Server
function Install-OpenSSH {
    Write-Host "Checking and installing OpenSSH components..."
    
    # Check OpenSSH Client
    $clientCapability = Get-WindowsCapability -Online | Where-Object Name -like "OpenSSH.Client*"
    if ($clientCapability.State -ne "Installed") {
        Write-Host "Installing OpenSSH Client..."
        try {
            Add-WindowsCapability -Online -Name OpenSSH.Client~~~~0.0.1.0
            Write-Host "OpenSSH Client installation completed"
        } catch {
            Write-Host "OpenSSH Client installation failed: $($_.Exception.Message)"
            return $false
        }
    } else {
        Write-Host "OpenSSH Client is already installed"
    }
    
    # Check OpenSSH Server
    $serverCapability = Get-WindowsCapability -Online | Where-Object Name -like "OpenSSH.Server*"
    if ($serverCapability.State -ne "Installed") {
        Write-Host "Installing OpenSSH Server..."
        try {
            Add-WindowsCapability -Online -Name OpenSSH.Server~~~~0.0.1.0
            Write-Host "OpenSSH Server installation completed"
        } catch {
            Write-Host "OpenSSH Server installation failed: $($_.Exception.Message)"
            return $false
        }
    } else {
        Write-Host "OpenSSH Server is already installed"
    }
    
    # Start and set SSH Agent service
    try {
        Set-Service -Name ssh-agent -StartupType Automatic
        Start-Service ssh-agent
        Write-Host "SSH Agent service started"
    } catch {
        Write-Host "SSH Agent service startup failed: $($_.Exception.Message)"
    }
    
    return $true
}

# Function: Generate SSH Key Pair
function New-SSHKeyPair {
    Write-Host "Checking SSH keys..."
    if (-not (Test-Path $keyPath)) {
        Write-Host "Generating new SSH key pair..."
        $sshDir = Split-Path $keyPath
        if (-not (Test-Path $sshDir)) {
            New-Item -ItemType Directory -Path $sshDir -Force | Out-Null
        }
        ssh-keygen -t rsa -b 4096 -f $keyPath -N ''
        Write-Host "SSH key pair generation completed"
    } else {
        Write-Host "SSH keys already exist"
    }
}

# Function: SSH with automatic password input using plink
function Invoke-SSHWithPassword {
    param(
        [string]$Server,
        [string]$Username,
        [string]$Password,
        [string]$Command,
        [int]$TimeoutSeconds = 30
    )
    
    # Try plink first
    $plinkPath = Get-Command plink -ErrorAction SilentlyContinue
    if ($plinkPath) {
        Write-Host "Using plink for automatic password input..."
        try {
            # Use plink with password parameter
            $result = & plink -ssh -batch -pw $Password $Username@$Server $Command 2>&1
            if ($LASTEXITCODE -eq 0) {
                Write-Host "plink command executed successfully"
                if ($result) { $result | ForEach-Object { Write-Host $_ } }
                return $true
            } else {
                Write-Host "plink command failed with exit code: $LASTEXITCODE"
                Write-Host "plink output: $result"
                return $false
            }
        } catch {
            Write-Host "plink failed: $($_.Exception.Message)"
            return $false
        }
    } else {
        Write-Host "plink not available, checking for installation..."
        
        # Try to install plink via winget
        try {
            Write-Host "Attempting to install PuTTY (includes plink) via winget..."
            $wingetResult = & winget install PuTTY.PuTTY --silent --accept-package-agreements --accept-source-agreements 2>&1
            if ($LASTEXITCODE -eq 0) {
                Write-Host "PuTTY installed successfully"
                # Refresh PATH and try again
                $env:PATH = [System.Environment]::GetEnvironmentVariable("PATH", "Machine") + ";" + [System.Environment]::GetEnvironmentVariable("PATH", "User")
                $plinkPath = Get-Command plink -ErrorAction SilentlyContinue
                if ($plinkPath) {
                    Write-Host "plink is now available, retrying..."
                    $result = & plink -ssh -batch -pw $Password $Username@$Server $Command 2>&1
                    if ($LASTEXITCODE -eq 0) {
                        if ($result) { $result | ForEach-Object { Write-Host $_ } }
                        return $true
                    }
                }
            } else {
                Write-Host "Failed to install PuTTY via winget: $wingetResult"
            }
        } catch {
            Write-Host "Error installing PuTTY: $($_.Exception.Message)"
        }
    }
    
    Write-Host "plink not available, falling back to manual configuration..."
    return $false
}

# Function: SSH with PowerShell Process and input redirection using expect-like approach
function Invoke-SSHWithProcessInput {
    param(
        [string]$Server,
        [string]$Username,
        [string]$Password,
        [string]$Command,
        [int]$TimeoutSeconds = 30
    )
    
    try {
        # Create an expect-like script for Windows
        $expectScript = @"
spawn ssh -o StrictHostKeyChecking=no -o PasswordAuthentication=yes $Username@$Server "$Command"
expect {
    "*password*" { send "$Password\r"; exp_continue }
    "*Password*" { send "$Password\r"; exp_continue }
    eof
}
"@
        $expectPath = "$env:TEMP\ssh_expect_$([System.Guid]::NewGuid().ToString('N').Substring(0,8)).exp"
        $expectScript | Out-File -FilePath $expectPath -Encoding ASCII
        
        # Try to use expect if available, otherwise use PowerShell native approach
        $expectExe = Get-Command "expect.exe" -ErrorAction SilentlyContinue
        if ($expectExe) {
            Write-Host "Using expect for SSH authentication..."
            $process = Start-Process -FilePath "expect.exe" -ArgumentList $expectPath -Wait -PassThru -NoNewWindow -RedirectStandardOutput "$env:TEMP\ssh_output.txt" -RedirectStandardError "$env:TEMP\ssh_error.txt"
        } else {
            Write-Host "Expect not available, using PowerShell native method..."
            # Use PowerShell's System.Diagnostics.Process with input simulation
            $psi = New-Object System.Diagnostics.ProcessStartInfo
            $psi.FileName = "ssh"
            $psi.Arguments = "-o StrictHostKeyChecking=no -o PasswordAuthentication=yes $Username@$Server `"$Command`""
            $psi.UseShellExecute = $false
            $psi.RedirectStandardInput = $true
            $psi.RedirectStandardOutput = $true
            $psi.RedirectStandardError = $true
            $psi.CreateNoWindow = $true
            
            $process = New-Object System.Diagnostics.Process
            $process.StartInfo = $psi
            $process.Start() | Out-Null
            
            # Wait a moment for SSH to prompt for password
            Start-Sleep -Milliseconds 1000
            
            # Send password multiple times to handle different prompts
            for ($i = 0; $i -lt 3; $i++) {
                try {
                    $process.StandardInput.WriteLine($Password)
                    Start-Sleep -Milliseconds 500
                } catch {
                    break
                }
            }
            
            $process.StandardInput.Close()
            $completed = $process.WaitForExit($TimeoutSeconds * 1000)
            
            if ($completed -and $process.ExitCode -eq 0) {
                $output = $process.StandardOutput.ReadToEnd()
                if ($output.Trim()) { Write-Host $output }
                return $true
            } else {
                $error = $process.StandardError.ReadToEnd()
                Write-Host "SSH command execution failed: $error"
                return $false
            }
        }
        
        # Clean up expect script
        Remove-Item $expectPath -Force -ErrorAction SilentlyContinue
        
        if ($process.ExitCode -eq 0) {
            $output = Get-Content "$env:TEMP\ssh_output.txt" -ErrorAction SilentlyContinue
            if ($output) { 
                $output | ForEach-Object { Write-Host $_ }
            }
            return $true
        } else {
            $error = Get-Content "$env:TEMP\ssh_error.txt" -ErrorAction SilentlyContinue
            Write-Host "SSH command execution failed: $error"
            return $false
        }
    } catch {
        Write-Host "SSH process exception: $($_.Exception.Message)"
        return $false
    } finally {
        # Clean up temporary files
        Remove-Item "$env:TEMP\ssh_output.txt" -Force -ErrorAction SilentlyContinue
        Remove-Item "$env:TEMP\ssh_error.txt" -Force -ErrorAction SilentlyContinue
        if ($process -and -not $process.HasExited) {
            $process.Kill()
        }
        if ($process) {
            $process.Dispose()
        }
    }
}

# Function: Install PuTTY (includes plink)
function Install-PuTTY {
    Write-Host "Checking and installing PuTTY..."
    
    # Check if plink is already available
    $plinkPath = Get-Command plink -ErrorAction SilentlyContinue
    if ($plinkPath) {
        Write-Host "plink is already available"
        return $true
    }
    
    # Try to install PuTTY using Chocolatey first
    try {
        $chocoPath = Get-Command choco -ErrorAction SilentlyContinue
        if ($chocoPath) {
            Write-Host "Installing PuTTY using Chocolatey..."
            $chocoResult = choco install putty -y 2>&1
            if ($LASTEXITCODE -eq 0) {
                Write-Host "PuTTY installed successfully via Chocolatey"
                # Refresh PATH
                $env:PATH = [System.Environment]::GetEnvironmentVariable("PATH", "Machine") + ";" + [System.Environment]::GetEnvironmentVariable("PATH", "User")
                $plinkPath = Get-Command plink -ErrorAction SilentlyContinue
                if ($plinkPath) {
                    return $true
                }
            }
        }
    } catch {
        Write-Host "Chocolatey installation failed, trying other methods"
    }
    
    # Try to install PuTTY using winget
    try {
        $wingetPath = Get-Command winget -ErrorAction SilentlyContinue
        if ($wingetPath) {
            Write-Host "Installing PuTTY using winget..."
            $wingetResult = winget install --id=PuTTY.PuTTY -e --silent 2>&1
            if ($LASTEXITCODE -eq 0) {
                Write-Host "PuTTY installed successfully via winget"
                # Refresh PATH
                $env:PATH = [System.Environment]::GetEnvironmentVariable("PATH", "Machine") + ";" + [System.Environment]::GetEnvironmentVariable("PATH", "User")
                $plinkPath = Get-Command plink -ErrorAction SilentlyContinue
                if ($plinkPath) {
                    return $true
                }
            }
        }
    } catch {
        Write-Host "winget installation failed"
    }
    
    Write-Host "Unable to automatically install PuTTY, please install manually"
    return $false
}

# Function: Configure SSH Key using plink for automatic configuration
function Set-SSHKeyWithPowerShell {
    param(
        [string]$Server,
        [string]$Username,
        [string]$Password,
        [string]$PublicKey
    )
    
    Write-Host "Configuring SSH key for $Server..."
    
    # Ensure PuTTY is installed
    if (-not (Install-PuTTY)) {
        Write-Host "PuTTY installation failed, falling back to manual configuration"
        return $false
    }
    
    # Use interactive SSH for manual password input
    try {
        Write-Host "Using interactive SSH (you will need to enter password manually)..."
        
        # Create .ssh directory
        Write-Host "Creating .ssh directory on $Server..."
        Write-Host "Please enter password when prompted:"
        $result1 = & ssh "$Username@$Server" "mkdir -p ~/.ssh; chmod 700 ~/.ssh" 2>&1
        if ($LASTEXITCODE -ne 0) {
            Write-Host "Failed to create .ssh directory: $result1"
            throw "Directory creation failed"
        }
        
        # Add public key to authorized_keys
        Write-Host "Adding public key to authorized_keys on $Server..."
        Write-Host "Please enter password when prompted:"
        $escapedKey = $PublicKey.Replace("'", "'\''")
        $result2 = & ssh "$Username@$Server" "echo '$escapedKey' >> ~/.ssh/authorized_keys; chmod 600 ~/.ssh/authorized_keys" 2>&1
        if ($LASTEXITCODE -ne 0) {
            Write-Host "Failed to add public key: $result2"
            throw "Key addition failed"
        }
        
        # Remove duplicate keys
        Write-Host "Cleaning duplicate keys on $Server..."
        Write-Host "Please enter password when prompted:"
        $result3 = & ssh "$Username@$Server" "sort ~/.ssh/authorized_keys | uniq > ~/.ssh/authorized_keys.tmp; mv ~/.ssh/authorized_keys.tmp ~/.ssh/authorized_keys" 2>&1
        if ($LASTEXITCODE -ne 0) {
            Write-Host "Failed to clean duplicate keys: $result3"
        }
        
        Write-Host "SSH key configuration successful: $Server"
        return $true
        
    } catch {
        Write-Host "SSH configuration failed: $($_.Exception.Message)"
        return $false
    }
    
    # If automatic configuration fails, provide manual instructions
    Write-Host "Automatic configuration failed. Please configure manually:"
    Write-Host "1. Copy the public key: $PublicKey"
    Write-Host "2. Connect to server: ssh $Username@$Server"
    Write-Host "3. Create .ssh directory: mkdir -p ~/.ssh && chmod 700 ~/.ssh"
    Write-Host "4. Add public key: echo '$PublicKey' >> ~/.ssh/authorized_keys"
    Write-Host "5. Set permissions: chmod 600 ~/.ssh/authorized_keys"
    Write-Host "6. Remove duplicates: sort ~/.ssh/authorized_keys | uniq > ~/.ssh/authorized_keys.tmp && mv ~/.ssh/authorized_keys.tmp ~/.ssh/authorized_keys"
    Write-Host ""
    
    # Return false for automatic processing without manual intervention
    return $false
}

# Function: Manual SSH Key Configuration
function Set-SSHKeyManual {
    param(
        [string]$server,
        [string]$username,
        [string]$password,
        [string]$publicKey
    )
    
    Write-Host "Manual configuration mode - $server"
    Write-Host "Please follow these steps for manual configuration:"
    Write-Host "1. Copy public key content: $publicKey"
    Write-Host "2. Connect to server: ssh $username@$server"
    Write-Host "3. Create directory: mkdir -p ~/.ssh; chmod 700 ~/.ssh"
    Write-Host "4. Add public key: echo '$publicKey' >> ~/.ssh/authorized_keys"
    Write-Host "5. Set permissions: chmod 600 ~/.ssh/authorized_keys"
    return $false
}

# Function: Clean known_hosts
function Clear-KnownHosts {
    $knownHostsPath = "$env:USERPROFILE\.ssh\known_hosts"
    if (Test-Path $knownHostsPath) {
        Write-Host "Cleaning known_hosts file..."
        foreach ($server in $servers) {
            ssh-keygen -R $server 2>$null
        }
    }
}

# Function: Test SSH Connection (key-based authentication)
function Test-SSHConnection {
    param(
        [string]$server, 
        [string]$username
    )
    
    Write-Host "Testing SSH key-based connection: $server"
    
    try {
        # Test SSH connection without password (key-based)
        $result = & ssh -o StrictHostKeyChecking=no -o PasswordAuthentication=no -o ConnectTimeout=10 $username@$server "echo 'SSH key-based connection successful'" 2>&1
        
        if ($LASTEXITCODE -eq 0) {
            Write-Host "SSH key-based connection test successful: $server"
            return $true
        } else {
            Write-Host "SSH key-based connection test failed: $server - $result"
            return $false
        }
    } catch {
        Write-Host "SSH connection test exception: $server - $($_.Exception.Message)"
        return $false
    }
}

# Main Program
Write-Host "=== SSH Key Auto Configuration Script ==="
Write-Host "Target servers: $($servers -join ', ')"
Write-Host "Username: $username"
Write-Host "Script version: 2.0 (Using PowerShell native methods)"
Write-Host ""

# 1. Install OpenSSH components
$opensshInstalled = Install-OpenSSH
if (-not $opensshInstalled) {
    Write-Host "OpenSSH installation failed, cannot continue"
    exit 1
}

# 2. Generate SSH key pair
New-SSHKeyPair

# 3. Read public key content
if (-not (Test-Path $pubKeyPath)) {
    Write-Host "Error: Public key file does not exist: $pubKeyPath"
    exit 1
}

$publicKeyContent = Get-Content $pubKeyPath -Raw
$publicKeyContent = $publicKeyContent.Trim()

# 4. Clean known_hosts
Clear-KnownHosts

# 5. Configure SSH keys
$successCount = 0
foreach ($server in $servers) {
    Write-Host ""
    Write-Host "--- Configuring server: $server ---"
    
    $result = Set-SSHKeyWithPowerShell -Server $server -Username $username -Password $password -PublicKey $publicKeyContent
    
    if ($result) {
        $successCount++
    } else {
        Write-Host "Server $server configuration failed, trying manual configuration method"
        Set-SSHKeyManual -server $server -username $username -password $password -publicKey $publicKeyContent
    }
}

# 6. Test connections
Write-Host ""
Write-Host "=== Testing SSH Connections ==="
$testSuccessCount = 0
foreach ($server in $servers) {
    if (Test-SSHConnection -server $server -username $username) {
        $testSuccessCount++
    }
}

# 7. Summary
Write-Host ""
Write-Host "=== Configuration Complete ==="
Write-Host "Successfully configured: $successCount/$($servers.Count) servers"
Write-Host "Connection tests: $testSuccessCount/$($servers.Count) servers"

if ($testSuccessCount -eq $servers.Count) {
    Write-Host "All servers SSH key configuration successful!"
    exit 0
} else {
    Write-Host "Some servers configuration failed, please check logs"
    exit 1
}